import numpy as np
import scipy.linalg
def mldivide_new(A,b):
    [mm,nn]=np.shape(A)
    ss=np.shape(b)[1]
    x=np.zeros((nn,ss))
    if np.shape(A)[0]==np.shape(A)[1]:
        if (A==np.tril(A)).all():
            for ii in range(ss):
                x[:,ii]=backward_sub(A,b[:,ii])
            return
        elif (A==np.triu(A)).all():
            for ii in range(ss):
                x[:,ii]=backward_sub(A,b[:,ii])
            return
        else:
            if (A==A.T).all():
                [R,p]=np.linalg.cholesky(A)
                if (p==0):
                    for ii in range(ss):
                        x[:,ii]=backward_sub(R,forward_sub(R.T,b[:,ii]))
                    return
        [L,U,P]=scipy.linalg.lu(A)
        for ii in range(ss):
            x[:,ii]=backward_sub(U,forward_sub(L,P.dot(b[:,ii])))
        return
    else:
        print(np.shape(A))
        [Q,R]=np.linalg.qr(A)
        Q=np.hstack((Q,np.array([[ii] for ii in Q[:,-1]])))
        R=np.vstack((R,R[-1,:]))
        Q[np.abs(Q)<=2.261143e-14]=0
        R[np.abs(R)<=2.261143e-14]=0
        print(np.shape(Q))
        print(np.shape(R))
        print(Q)
        print("R=", R)
        for jj in range(ss):
            print(np.shape(x))
            x1 = backward_sub(R, Q.T.dot(b[:, jj]))
            x1=[i[0] for i in x1]
            x[:,jj]=x1
        return x
    return x
def forward_sub(A,b):
    [m,n]=np.shape(A)
    x=np.zeros((n,1))
    for i in range(0,n):

        if np.abs(A[i,i])>=2.261143e-14:
            x[i]=(b[i]-A[i,0:i-1].dot(x[0,i-1]))/A[i,i]

        elif b[i]-A[i,0:i-1].dot(x[0:i-1])==0:
            x[i]=0

    return x
def backward_sub(A,b):
    [m, n] = np.shape(A)
    x = np.zeros((n, 1))
    for i in range(n-1,-1,-1):
        print(int(A[i,i]))
        print(type(A[i,i]))
        if int(A[i,i])!=0:
            x[i]=(b[i]-A[i,i+1:-1].dot(x[i+1:-1]))/A[i,i]
        elif b[i]-A[i,i+1:-1].dot(x[i+1:-1])==0:
            x[i]=0

    return x


if __name__=='__main__':
    A=np.matrix([[8,-3,2],[4,11,-1],[6,3,12]])
#    B=np.array([1,2,3])   
    B=np.matrix([1,2,2]).T
    x=mldivide_new(A,B)
    print(x)
    x2 = np.linalg.inv(A).dot(B)
    print(x2)
    
